`timescale 1ns/1ps
`include "../define.sv"

module systolic_array_tb;

// 时钟和复位信号
reg clk;
reg reset;

// 矩阵输入接口
reg [31:0] a_din [0:31];
reg a_wr_en [0:31];
reg [31:0] b_din [0:31];
reg b_wr_en [0:31];
reg [31:0] c_din [0:31][0:31];
reg c_wr_en [0:31][0:31];

// 控制信号
reg [1:0] matrix_type;
reg [1:0] data_type;
reg mix_precision;

// 输出接口
wire [31:0] c_out_muxed;

// 预期结果存储
reg [31:0] expected_D [0:31][0:31];

//--------------------- 类型定义 ---------------------
typedef struct {
    logic [7:0] a_data [0:31][0:15];  // A矩阵（按行存储，int8或int4）
    logic [7:0] b_data [0:15][0:31];  // B矩阵（按列存储，int8或int4）
    logic [31:0] c_data [0:31][0:31]; // C矩阵（大端int32）
    logic [31:0] d_data [0:31][0:31]; // 预期结果
    int          m_size;              // A矩阵行数
    int          n_size;              // B矩阵列数
    int          k_size;              // 内积维度
    string       test_name;
    string       data_path;           // 数据文件路径
    bit          is_int4;             // 是否为int4数据类型
} matrix_test_case;

matrix_test_case test_queue[$]; // 测试用例队列

// 实例化被测模块
systolic_system uut (
    .clk(clk),
    .reset(reset),
    .a_din(a_din),
    .a_wr_en(a_wr_en),
    .b_din(b_din),
    .b_wr_en(b_wr_en),
    .c_din(c_din),
    .c_wr_en(c_wr_en),
    .matrix_type(matrix_type),
    .data_type(data_type),
    .mix_precision(mix_precision),
    .c_out_muxed(c_out_muxed)
);

// 时钟生成
initial begin
    clk = 0;
    forever #2.5 clk = ~clk;
end

initial begin
    // 初始化
    int wait_count; // 声明移到foreach外部
    reset = 1;
    matrix_type = `m16n16k16;
    data_type = 2'b01;
    mix_precision = 0;
    foreach(a_wr_en[i]) a_wr_en[i] = 0;
    foreach(b_wr_en[i]) b_wr_en[i] = 0;
    foreach(c_wr_en[i,j]) c_wr_en[i][j] = 0;

    // 加载测试数据到队列
    load_test_cases();

    // 释放复位并执行测试用例
    #30 reset = 0;

    foreach(test_queue[i]) begin
        $display("start test case: %s", test_queue[i].test_name);
        
        // 根据矩阵尺寸设置matrix_type
        set_matrix_type(test_queue[i]);
        
        // 根据数据类型设置data_type
        if (test_queue[i].is_int4) begin
            data_type = `INT4; // int4
        end else begin
            data_type = `INT8; // int8
        end
        $display("Set data_type = %d for %s", data_type, test_queue[i].is_int4 ? "INT4" : "INT8");
        
        write_fifo_data_generic(test_queue[i]);
        
        // 等待状态机进入OUTPUT状态，开始捕获数据
        $display("[%t] Waiting for array state to enter OUTPUT...", $time);
        wait_count = 0;
        while (uut.array_state != 2'b10 && wait_count < 1000) begin  // 等待OUTPUT状态
            @(posedge clk);
            wait_count++;
            if (wait_count % 100 == 0) begin
                $display("[%t] Still waiting for OUTPUT state, current state: %d, wait_count: %d", 
                         $time, uut.array_state, wait_count);
            end
        end
        
        if (wait_count >= 1000) begin
            $error("[%t] Timeout waiting for OUTPUT state! Current state: %d", $time, uut.array_state);
            continue; // 跳过这个测试用例
        end
        
        $display("[%t] Array state is now OUTPUT, starting data capture", $time);
        
        // 立即开始验证结果（在下一个时钟周期）
        @(posedge clk);
        verify_results_generic(test_queue[i]);
        
        // 等待PE复位完成后再开始下一个测试
        wait_for_pe_reset_complete();
    end
    
    $display("all test cases done!");
    $stop;
end

// 添加状态监控
initial begin
    automatic int last_state = -1;
    automatic int state_change_count = 0;
    
    forever begin
        @(posedge clk);
        if (uut.array_state != last_state) begin
            $display("[%t] Array State Change: %d -> %d (change #%d)", 
                     $time, last_state, uut.array_state, state_change_count);
            last_state = uut.array_state;
            state_change_count++;
        end
    end
end

//---------------------------------------------------------------------
// 任务：设置矩阵类型
//---------------------------------------------------------------------
task set_matrix_type(input matrix_test_case tc);
    case({tc.m_size[7:0], tc.n_size[7:0], tc.k_size[7:0]})
        {8'd16, 8'd16, 8'd16}: matrix_type = `m16n16k16;
        {8'd8, 8'd32, 8'd16}:  matrix_type = `m8n32k16;
        {8'd32, 8'd8, 8'd16}:  matrix_type = `m32n8k16;
        default: begin
            $error("Unsupported matrix size: %dx%dx%d", tc.m_size, tc.n_size, tc.k_size);
            matrix_type = `m16n16k16; // 默认值
        end
    endcase
    $display("Set matrix_type = %d for %dx%dx%d", matrix_type, tc.m_size, tc.n_size, tc.k_size);
endtask

//---------------------------------------------------------------------
// 任务：加载测试用例
//---------------------------------------------------------------------
task load_test_cases();
    matrix_test_case current_case;
    
    // 测试用例1: m16n16k16 int8
    current_case.test_name = "m16n16k16_int8_test";
    current_case.m_size = 16;
    current_case.n_size = 16;
    current_case.k_size = 16;
    current_case.is_int4 = 0;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int8_int32/m16n16k16/";
    load_single_test_case(current_case);
    
    // 测试用例2: m8n32k16 int8
    current_case.test_name = "m8n32k16_int8_test";
    current_case.m_size = 8;
    current_case.n_size = 32;
    current_case.k_size = 16;
    current_case.is_int4 = 0;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int8_int32/m8n32k16/";
    load_single_test_case(current_case);
    
    // 测试用例3: m32n8k16 int8
    current_case.test_name = "m32n8k16_int8_test";
    current_case.m_size = 32;
    current_case.n_size = 8;
    current_case.k_size = 16;
    current_case.is_int4 = 0;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int8_int32/m32n8k16/";
    load_single_test_case(current_case);
    
    // 测试用例4: m16n16k16 int4
    current_case.test_name = "m16n16k16_int4_test";
    current_case.m_size = 16;
    current_case.n_size = 16;
    current_case.k_size = 16;
    current_case.is_int4 = 1;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int4_int32/m16n16k16/";
    load_single_test_case(current_case);
    
    // 测试用例5: m8n32k16 int4
    current_case.test_name = "m8n32k16_int4_test";
    current_case.m_size = 8;
    current_case.n_size = 32;
    current_case.k_size = 16;
    current_case.is_int4 = 1;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int4_int32/m8n32k16/";
    load_single_test_case(current_case);
    
    // 测试用例6: m32n8k16 int4
    current_case.test_name = "m32n8k16_int4_test";
    current_case.m_size = 32;
    current_case.n_size = 8;
    current_case.k_size = 16;
    current_case.is_int4 = 1;
    current_case.data_path = "D:/FPGA/Full/TPU/SA/TPU/top/systolic_array/testbench/int4_int32/m32n8k16/";
    load_single_test_case(current_case);
    
    $display("Loaded %d test cases", test_queue.size());
endtask

//---------------------------------------------------------------------
// 任务：加载单个测试用例
//---------------------------------------------------------------------
task load_single_test_case(input matrix_test_case tc);
    automatic int fd;
    automatic string line;
    automatic string a_filename, b_filename, d_filename;
    automatic matrix_test_case current_case = tc;
    
    // 加载C矩阵（int32大端）
    fd = $fopen({current_case.data_path, "c_int32_big_endian_m", 
                 $sformatf("%0d", current_case.m_size), "n", 
                 $sformatf("%0d", current_case.n_size), "k", 
                 $sformatf("%0d", current_case.k_size), "_hex.txt"}, "r");
    if (fd) begin
        for (int i=0; i<current_case.m_size; i++) begin
            for (int j=0; j<current_case.n_size; j++) begin
                automatic logic [31:0] val;
                if ($fscanf(fd, "%8h", val) != 1) begin
                    $error("C matrix format error @(%0d,%0d)", i, j);
                end
                current_case.c_data[i][j] = val;
            end
        end
        $fclose(fd);
    end else $error("cannot open C matrix file for %s", current_case.test_name);
    
    // 加载A矩阵（INT8/INT4行主序）
    if (current_case.is_int4) begin
        a_filename = {current_case.data_path, "a_int4_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex.txt"};
    end else begin
        a_filename = {current_case.data_path, "a_int8_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex.txt"};
    end
    
    fd = $fopen(a_filename, "r");
    if (fd) begin
        for (int row=0; row<current_case.m_size; row++) begin
            for (int col=0; col<current_case.k_size; col++) begin
                automatic logic [7:0] val;
                if (current_case.is_int4) begin
                    // int4: 读取4bit十六进制，转换为8bit
                    automatic logic [3:0] val_4bit;
                    if ($fscanf(fd, "%1h", val_4bit) != 1) begin
                        $error("A matrix int4 format error @(%0d,%0d)", row, col);
                    end
                    // 处理int4有符号补码
                    if (val_4bit >= 4'h8) begin
                        val = {4'hF, val_4bit}; // 负数扩展
                    end else begin
                        val = {4'h0, val_4bit}; // 正数扩展
                    end
                end else begin
                    // int8: 直接读取8bit十六进制
                    if ($fscanf(fd, "%2h", val) != 1) begin
                        $error("A matrix int8 format error @(%0d,%0d)", row, col);
                    end
                end
                current_case.a_data[row][col] = val;
            end
        end
        $fclose(fd);
    end else $error("cannot open A matrix file for %s", current_case.test_name);

    // 加载B矩阵（INT8/INT4，按正常行优先顺序读取）
    if (current_case.is_int4) begin
        b_filename = {current_case.data_path, "b_int4_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex.txt"};
    end else begin
        b_filename = {current_case.data_path, "b_int8_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex.txt"};
    end
    
    fd = $fopen(b_filename, "r");
    if (fd) begin
        // 按正常行优先顺序读取数据
        for (int row=0; row<current_case.k_size; row++) begin
            for (int col=0; col<current_case.n_size; col++) begin
                automatic logic [7:0] val;
                
                if (current_case.is_int4) begin
                    // int4: 读取4bit十六进制，转换为8bit
                    automatic logic [3:0] val_4bit;
                    if ($fscanf(fd, "%1h", val_4bit) != 1) begin
                        $error("B matrix int4 format error @(%0d,%0d)", row, col);
                    end
                    // 处理int4有符号补码
                    if (val_4bit >= 4'h8) begin
                        val = {4'hF, val_4bit}; // 负数扩展
                    end else begin
                        val = {4'h0, val_4bit}; // 正数扩展
                    end
                end else begin
                    // int8: 直接读取8bit十六进制
                    if ($fscanf(fd, "%2h", val) != 1) begin
                        $error("B矩阵 int8 format error @(%0d,%0d)", row, col);
                    end
                end
                
                current_case.b_data[row][col] = val;
            end
        end
        $fclose(fd);
    end else $error("cannot open B matrix file for %s", current_case.test_name);

    // 加载预期结果
    if (current_case.is_int4) begin
        d_filename = {current_case.data_path, "d_int32_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex_new.txt"};
    end else begin
        d_filename = {current_case.data_path, "d_int32_m", 
                     $sformatf("%0d", current_case.m_size), "n", 
                     $sformatf("%0d", current_case.n_size), "k", 
                     $sformatf("%0d", current_case.k_size), "_hex.txt"};
    end
    
    fd = $fopen(d_filename, "r");
    if (fd) begin
        for (int i=0; i<current_case.m_size; i++) begin
            for (int j=0; j<current_case.n_size; j++) begin
                if ($fscanf(fd, "%8h", current_case.d_data[i][j]) != 1) begin
                    $error("D matrix format error @(%0d,%0d)", i, j);
                end
            end
        end
        $fclose(fd);
    end else $error("cannot open D matrix file for %s", current_case.test_name);

    test_queue.push_back(current_case);
    $display("success load test case: %s (%dx%dx%d, %s)", 
             current_case.test_name, current_case.m_size, current_case.n_size, current_case.k_size,
             current_case.is_int4 ? "INT4" : "INT8");
endtask

//---------------------------------------------------------------------
// 任务：写入FIFO数据（通用版本）
//---------------------------------------------------------------------
task write_fifo_data_generic(input matrix_test_case tc);
    automatic int row, col;
    automatic int a_cycles, b_cycles, max_cycles;
begin
    // 初始化C矩阵输入
    for(int i=0; i<32; i++) begin
        for(int j=0; j<32; j++) begin
            c_din[i][j] = 32'h0;
        end
    end

     // 写入C矩阵（根据matrix_type映射）
    for(int cycle=0; cycle<tc.m_size*tc.n_size; cycle++) begin
        @(posedge clk);
        // 默认全部清零
        foreach(c_wr_en[i,j]) c_wr_en[i][j] = 0;
        foreach(c_din[i,j]) c_din[i][j] = 0;

        case (matrix_type)
            `m16n16k16: begin
                row = cycle / 16;
                col = cycle % 16;
                c_wr_en[row][col] = 1'b1;
                c_din[row][col] = tc.c_data[row][col];
            end
            `m8n32k16: begin
                row = cycle / 32;
                col = cycle % 32;
                if (col < 16) begin
                    // 左半部分，写到[0:7][0:15]
                    c_wr_en[row][col] = 1'b1;
                    c_din[row][col] = tc.c_data[row][col];
                end else begin
                    // 右半部分，写到[8:15][col-16]
                    c_wr_en[row+8][col-16] = 1'b1;
                    c_din[row+8][col-16] = tc.c_data[row][col];
                end
            end
            `m32n8k16: begin
                row = cycle / 8;
                col = cycle % 8;
                if (row < 16) begin
                    // 上半部分，写到[0:15][0:7]
                    c_wr_en[row][col] = 1'b1;
                    c_din[row][col] = tc.c_data[row][col];
                end else begin
                    // 下半部分，写到[0:15][col+8]
                    c_wr_en[row-16][col+8] = 1'b1;
                    c_din[row-16][col+8] = tc.c_data[row][col];
                end
            end
            default: begin
                // 默认不写
            end
        endcase
        #1;
    end
    
    @(posedge clk);
    foreach(c_wr_en[i,j]) c_wr_en[i][j] = 0;
    
    // 等待几个时钟周期，确保C矩阵写入完全完成
    repeat(5) @(posedge clk);

    // 打印调试信息
    $display("Starting AB matrix write for %s...", tc.test_name);
    $display("Matrix dimensions: %dx%dx%d (%s)", tc.m_size, tc.n_size, tc.k_size, tc.is_int4 ? "INT4" : "INT8");
    $display("A matrix first row: %h %h %h %h", tc.a_data[0][0], tc.a_data[0][1], tc.a_data[0][2], tc.a_data[0][3]);
    $display("B matrix first col: %h %h %h %h", tc.b_data[0][0], tc.b_data[1][0], tc.b_data[2][0], tc.b_data[3][0]);
    $display("B matrix last col: %h %h %h %h", tc.b_data[0][tc.n_size-1], tc.b_data[1][tc.n_size-1], tc.b_data[2][tc.n_size-1], tc.b_data[3][tc.n_size-1]);

    // 计算A和B各自需要的写入周期
    a_cycles = tc.m_size + tc.k_size - 1;
    b_cycles = tc.n_size + tc.k_size - 1;
    max_cycles = (a_cycles > b_cycles) ? a_cycles : b_cycles;
    
    $display("Matrix write cycles: A=%d, B=%d, max=%d", a_cycles, b_cycles, max_cycles);
    
    // 写入A和B矩阵
    for (int cycle=0; cycle<max_cycles; cycle++) begin
        @(posedge clk);
        
        // 初始化所有通道为0
        for (int i=0; i<32; i++) begin
            a_wr_en[i] = 1'b0;
            b_wr_en[i] = 1'b0;
            a_din[i]   = 32'h0;
            b_din[i]   = 32'h0;
        end
        
        // 写入A矩阵：每个FIFO在其活动窗口内写入
        for (int i=0; i<tc.m_size; i++) begin
            if (cycle >= i && (cycle-i) < tc.k_size) begin
                automatic int data_index = cycle - i;
                a_wr_en[i] = 1'b1;
                a_din[i]   = {24'h0, tc.a_data[i][data_index]}; 
            end
        end
        
        // 写入B矩阵：每个FIFO在其活动窗口内写入
        for (int j=0; j<tc.n_size; j++) begin
            if (cycle >= j && cycle < j + tc.k_size) begin
                automatic int data_index = cycle - j;
                
                b_wr_en[j] = 1'b1;
                b_din[j]   = {24'h0, tc.b_data[data_index][j]};
            end
        end
        
        // 打印调试信息（仅打印前几个周期和关键周期）
        if (cycle < 4 || cycle == tc.n_size-1 || cycle == max_cycles-1) begin
            $display("Cycle %d: Writing A/B data", cycle);
            for (int i=0; i<tc.m_size; i++) begin
                if (a_wr_en[i]) begin
                    $display("  A[%d] = %h", i, a_din[i]);
                end
            end
            for (int j=0; j<tc.n_size; j++) begin
                if (b_wr_en[j]) begin
                    $display("  B[%d] = %h", j, b_din[j]);
                end
            end
        end
    end
    
    // 关闭A和B的写使能
    @(posedge clk);
    for (int i=0; i<32; i++) begin
        a_wr_en[i] = 0;
        b_wr_en[i] = 0;
    end
    
    // 等待几个时钟周期，确保所有写入完成
    repeat(5) @(posedge clk);
    
    // 验证B矩阵FIFO写入情况
    //verify_b_fifo_write_sequence(tc);
end
endtask


//---------------------------------------------------------------------
// 任务：验证结果（通用版本）
//---------------------------------------------------------------------
task verify_results_generic(input matrix_test_case tc);
    reg [31:0] captured_c_out [0:31][0:31];
    automatic integer cycle_count = 0;
    automatic bit error_flag = 0;
    automatic int expected_cycles = tc.m_size * tc.n_size;
    automatic int timeout_count = 0;

    $display("[%t] start verify test case: %s (%dx%dx%d, %s)", 
             $time, tc.test_name, tc.m_size, tc.n_size, tc.k_size, tc.is_int4 ? "INT4" : "INT8");
    
    // 初始化捕获数组
    for (int i = 0; i < 32; i++) begin
        for (int j = 0; j < 32; j++) begin
            captured_c_out[i][j] = 32'd0;
        end
    end
    
    // 捕获输出：当阵列处于OUTPUT状态时，每周期获取c_out_muxed信号
    fork
        begin
            // 持续捕获数据直到完成
            while (cycle_count < expected_cycles) begin
                @(posedge clk);
                
                // 检查阵列是否处于OUTPUT状态
                if (uut.array_state == 2'b10) begin
                    // 直接获取c_out_muxed信号的值
                    captured_c_out[cycle_count / tc.n_size][cycle_count % tc.n_size] = c_out_muxed;
                    cycle_count++;
                    
                    // 打印调试信息（前几个周期和最后几个周期）
                    if (cycle_count <= 4 || cycle_count > expected_cycles - 4) begin
                        $display("[%t] Cycle %d: captured[%d][%d] = %h", 
                                 $time, cycle_count, cycle_count / tc.n_size, cycle_count % tc.n_size, c_out_muxed);
                    end
                    
                    // 每50个周期打印一次进度
                    if (cycle_count % 50 == 0) begin
                        $display("[%t] Progress: %d/%d", $time, cycle_count, expected_cycles);
                    end
                end else begin
                    // 如果不在OUTPUT状态，打印状态信息
                    if (cycle_count > 0 && cycle_count % 100 == 0) begin
                        $display("[%t] Waiting for OUTPUT state, current state: %d, captured: %d/%d", 
                                 $time, uut.array_state, cycle_count, expected_cycles);
                    end
                end
            end
            
            $display("[%t] Successfully captured all %d outputs", $time, expected_cycles);
        end
        // 超时保护
        #50000 begin
            $error("[%s] timeout! not capture all output (captured %d/%d)", 
                  tc.test_name, cycle_count, expected_cycles);
            $display("[%t] Current array state: %d", $time, uut.array_state);
            disable fork;
        end
    join



    /* // 打印捕获的前几个结果用于调试
    $display("=== Captured Results (first few) ===");
    for (int i=0; i<((4 < tc.m_size) ? 4 : tc.m_size); i++) begin
        for (int j=0; j<((4 < tc.n_size) ? 4 : tc.n_size); j++) begin
            $display("PE[%2d][%2d] = %h", i, j, captured_c_out[i][j]);
        end
    end
    
    $display("=== Expected Results (first few) ===");
    for (int i=0; i<((4 < tc.m_size) ? 4 : tc.m_size); i++) begin
        for (int j=0; j<((4 < tc.n_size) ? 4 : tc.n_size); j++) begin
            $display("PE[%2d][%2d] = %h", i, j, tc.d_data[i][j]);
        end
    end */

    // 与预期结果比较
    for (int i=0; i<tc.m_size; i++) begin
        for (int j=0; j<tc.n_size; j++) begin
            if (captured_c_out[i][j] !== tc.d_data[i][j]) begin
                error_flag = 1;
                $error("[%s] PE[%2d][%2d] error: expected=%h, actual=%h", 
                      tc.test_name, i, j, tc.d_data[i][j], captured_c_out[i][j]);
            end
        end
    end

    if (!error_flag) begin
        $display("[%t] test case pass: %s", $time, tc.test_name);
    end else begin
        $display("[%t] test case fail: %s", $time, tc.test_name);
    end
endtask

//---------------------------------------------------------------------
// 任务：等待计算完成
//---------------------------------------------------------------------
task wait_for_calculation_complete(input matrix_test_case tc);
    automatic int timeout_count = 0;
    automatic int max_timeout = 5000; // 增加超时时间
    
    $display("[%t] Waiting for calculation to complete...", $time);
    
    // 等待最后一个PE计算完成
    while (!uut.calc_done[15][15] && timeout_count < max_timeout) begin
        @(posedge clk);
        timeout_count++;
        
        // 每1000个周期打印一次状态
        if (timeout_count % 1000 == 0) begin
            $display("[%t] Still waiting... timeout_count=%d, calc_done[15][15]=%b", 
                     $time, timeout_count, uut.calc_done[15][15]);
        end
    end
    
    if (timeout_count >= max_timeout) begin
        $error("[%t] Calculation timeout! calc_done[15][15]=%b", $time, uut.calc_done[15][15]);
    end else begin
        $display("[%t] Calculation completed! counter=%d", $time, uut.counter);
    end
    
    // 等待输出完成
    timeout_count = 0;
    while (!uut.output_complete && timeout_count < 1000) begin
        @(posedge clk);
        timeout_count++;
    end
    
    if (timeout_count >= 1000) begin
        $error("[%t] Output completion timeout!", $time);
    end else begin
        $display("[%t] Output completed!", $time);
    end
endtask

//---------------------------------------------------------------------
// 任务：等待PE复位完成
//---------------------------------------------------------------------
task wait_for_pe_reset_complete();
    automatic int wait_count = 0;
    
    // 等待状态机回到IDLE状态
    $display("[%t] Waiting for array state to return to IDLE...", $time);
    while (uut.array_state != 2'b00 && wait_count < 200) begin  // 等待IDLE状态
        @(posedge clk);
        wait_count++;
    end
    
    if (wait_count >= 200) begin
        $error("[%t] Array state transition timeout! Current state: %d", $time, uut.array_state);
    end else begin
        $display("[%t] Array state returned to IDLE after %d cycles", $time, wait_count);
    end
    
    // 额外等待几个周期确保状态稳定
    repeat(5) @(posedge clk);
endtask



endmodule 